{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "I8JvIOhzJqkG"
   },
   "source": [
    "# Computing Galactic Orbits of Stars with Gala"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fBgeJRwZBAqf"
   },
   "source": [
    "## Authors\n",
    "Adrian Price-Whelan, Stephanie T. Douglas\n",
    "\n",
    "## Learning Goals\n",
    "* Query the Gaia data release 2 catalog to retrieve data for a sample of well-measured, nearby stars\n",
    "* Define high-mass and low-mass stellar samples using color-magnitude selections\n",
    "* Calculate orbits of the high-mass and low-mass stars within the Galaxy to show that the typically younger stars (high-mass) have smaller vertical excursions\n",
    "\n",
    "## Keywords\n",
    "coordinates, astroquery, gala, galactic dynamics, astrometry, matplotlib, scatter plot, histogram\n",
    "\n",
    "\n",
    "## Companion Content\n",
    "Astropy Docs: [Description of the Galactocentric frame in astropy coordinates](\n",
    "http://docs.astropy.org/en/latest/generated/examples/coordinates/plot_galactocentric-frame.html#sphx-glr-generated-examples-coordinates-plot-galactocentric-frame-py)\n",
    "\n",
    "## Summary\n",
    "\n",
    "We will use data from the [Gaia mission](https://www.cosmos.esa.int/web/gaia) to get sky positions, distances (parallaxes), proper motions, and radial velocities for a set of stars that are close to the Sun. We will then transform these observed, heliocentric kinematic measurements to Galactocentric Cartesian coordinates and use the positions and velocities as initial conditions to compute the orbits of these stars in the galaxy using the [gala](http://gala.adrian.pw) Python package. We will compare the orbits of high-mass main sequence (i.e. young) stars to the orbits of lower-mass main sequence stars to show that young stars have smaller vertical amplitudes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rcN6qYxDGZch"
   },
   "source": [
    "## Installing Dependencies\n",
    "\n",
    "This tutorial depends on the Astropy affiliated packages `gala` and `astroquery`. Both of these packages can be pip-installed with:\n",
    "\n",
    "    pip install astro-gala astroquery"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FPN267eAJ3PE"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xoqHIixqBAqg"
   },
   "outputs": [],
   "source": [
    "# astropy imports\n",
    "import astropy.coordinates as coord\n",
    "from astropy.table import QTable\n",
    "import astropy.units as u\n",
    "from astroquery.gaia import Gaia\n",
    "\n",
    "# Third-party imports\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "\n",
    "# gala imports\n",
    "import gala.coordinates as gc\n",
    "import gala.dynamics as gd\n",
    "import gala.potential as gp\n",
    "from gala.units import galactic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sNPUtNHptm1R"
   },
   "source": [
    "## Scientific Background\n",
    "\n",
    "[The Gaia mission](https://www.cosmos.esa.int/web/gaia/) is an ESA mission that aims to measure the 3D positions and velocities of a large number of stars throughout the Milky Way. The primary mission objective is to enable studying the formation, structure, and evolutionary history of our Galaxy by measuring astrometry (sky position, parallax, and proper motion) for about 2 billion stars brighter than the Gaia $G$-band photometric magnitude $G \\lesssim 21$. By end of mission (~2022), Gaia will also provide multi-band photometry and low-resolution spectrophotometry for these sources, along with radial or line-of-sight velocities for a subsample of about 100 million stars.\n",
    "\n",
    "In April 2018, Gaia publicly released its first major catalog of data — data release 2 or DR2 — which provides a subset of these data to anyone with an internet connection. In this tutorial, we will use astrometry, radial velocities, and photometry for a small subset of DR2 to study the kinematics of different types of stars in the Milky Way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "45SP61VoBAqj"
   },
   "source": [
    "## Using `astroquery` to retrieve Gaia data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xsuf3f63BAqj"
   },
   "source": [
    "We'll start by querying the [Gaia science archive](http://gea.esac.esa.int/archive/) to download astrometric and kinematic data (parallax, proper motion, radial velocity) for a sample of stars near the Sun. We'll use data exclusively from [data release 2 (DR2)](https://www.cosmos.esa.int/web/gaia/data-release-2) from the *Gaia* mission. For the demonstration here, let's grab data for a random subset of 4096 stars within a distance of 100 pc from the Sun that have high signal-to-noise astrometric measurements.\n",
    "\n",
    "To perform the query and to retrieve the data, we'll use the *Gaia* module in the [astroquery](https://astroquery.readthedocs.io/en/latest/gaia/gaia.html) package, `astroquery.gaia`. This module expects us to provide an SQL query to select the data we want (technically it should be an [ADQL](https://gea.esac.esa.int/archive-help/adql/index.html) query, which is similar to SQL but provides some additional functionality for astronomy; to learn more about ADQL syntax and options, [this guide](https://www.gaia.ac.uk/data/gaia-data-release-1/adql-cookbook) provides an introduction). We don't need all of the columns that are available in DR2, so we'll limit our query to request the sky position (`ra`, `dec`), parallax, proper motion components (`pmra`, `pmdec`), radial velocity, and magnitudes (`phot_*_mean_mag`). More information about the available columns is in the [Gaia DR2 data model](https://gea.esac.esa.int/archive/documentation/GDR2/Gaia_archive/chap_datamodel/sec_dm_main_tables/ssec_dm_gaia_source.html). \n",
    "\n",
    "To select stars that have high signal-to-noise parallaxes, we'll use the filter ``parallax_over_error > 10`` to select stars that have small fractional uncertainties. We'll also use the filter ``radial_velocity IS NOT null`` to only select stars that have measured radial velocities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BHJATifiBAqk"
   },
   "outputs": [],
   "source": [
    "query_text = '''SELECT TOP 4096 ra, dec, parallax, pmra, pmdec, radial_velocity,\n",
    "phot_g_mean_mag, phot_bp_mean_mag, phot_rp_mean_mag\n",
    "FROM gaiadr2.gaia_source\n",
    "WHERE parallax_over_error > 10 AND\n",
    "    parallax > 10 AND\n",
    "    radial_velocity IS NOT null\n",
    "ORDER BY random_index\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tDtPL9waBAqm"
   },
   "source": [
    "We now pass this query to the `Gaia.launch_job()` class method to create an anonymous job in the `Gaia` science archive to run our query. To retrieve the results of this query as an Astropy `Table` object, we then use the `job.get_results()` method. Note that you may receive a number of warnings (output lines that begin with ``WARNING:``) from the ``astropy.io.votable`` package — these are expected, and it's OK to ignore these warnings (the `Gaia` archive returns a slightly invalid VOTable)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5EbjfMrwBAqn"
   },
   "outputs": [],
   "source": [
    "# Note: the following lines require an internet connection, so we have \n",
    "# provided the results of this query as a FITS file included with the \n",
    "# tutorials repository. If you have an internet connection, feel free\n",
    "# to uncomment these lines to retrieve the data with `astroquery`:\n",
    "# job = Gaia.launch_job(query_text)\n",
    "# gaia_data = job.get_results()\n",
    "# gaia_data.write('gaia_data.fits')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CUcxP-lNGZct"
   },
   "outputs": [],
   "source": [
    "gaia_data = QTable.read('gaia_data.fits')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "76S58ivIBAqs"
   },
   "source": [
    "The `data` object is now an Astropy `Table` called `gaia_data` that contains `Gaia` data for 4096 random stars within 100 pc (or with a parallax > 10 mas) of the Sun, as we requested. Let's look at the first four rows of the table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dYMghBSnBAqu"
   },
   "outputs": [],
   "source": [
    "gaia_data[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fyx_IWdxBAq0"
   },
   "source": [
    "Note that the table columns already contain units! They are indicated in the second row of the header."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ARy1srbkBAq1"
   },
   "source": [
    "## Using `astropy.coordinates` to represent and transform stellar positions and velocities\n",
    "\n",
    "Let's double check that the farthest star is still within 100 pc, as we expect from the parallax selection we did in the query above. To do this, we'll create an Astropy `Distance` object using the parallax (*Note: this inverts the parallax to compute the distance! This is only a good approximation when the parallax signal to noise is large, as we ensured in the query above with `parallax_over_error > 10`*):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "luS44AQzBAq2"
   },
   "outputs": [],
   "source": [
    "dist = coord.Distance(parallax=u.Quantity(gaia_data['parallax']))\n",
    "dist.min(), dist.max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SEx8IC9PBAq8"
   },
   "source": [
    "It looks like the closest star in our sample is about 9 pc away, and the farthest is almost 100 pc, as we expected.\n",
    "\n",
    "We next want to convert the coordinate position and velocity data from heliocentric, spherical values to Galactocentric, Cartesian values. We'll do this using the [Astropy coordinates](http://docs.astropy.org/en/latest/coordinates/index.html) transformation machinery. To make use of this functionality, we first have to create a `SkyCoord` object from the `Gaia` data we downloaded. The `Gaia` DR2 data are in the ICRS (equatorial) reference frame, which is also the default frame when creating new `SkyCoord` objects, so we don't need to specify the frame below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vslMn05uBAq9"
   },
   "outputs": [],
   "source": [
    "c = coord.SkyCoord(ra=gaia_data['ra'], \n",
    "                   dec=gaia_data['dec'],\n",
    "                   distance=dist,\n",
    "                   pm_ra_cosdec=gaia_data['pmra'], \n",
    "                   pm_dec=gaia_data['pmdec'],\n",
    "                   radial_velocity=gaia_data['radial_velocity'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_jeqHDQYBArA"
   },
   "source": [
    "Note: as described in the [Gaia DR2 data model](https://gea.esac.esa.int/archive/documentation/GDR2/Gaia_archive/chap_datamodel/sec_dm_main_tables/ssec_dm_gaia_source.html), the Gaia column `pmra` contains the cos(dec) term. In Astropy coordinates, the name of this component is `pm_ra_cosdec`.\n",
    "\n",
    "Let's again look at the first four coordinates in the `SkyCoord` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "r7fNULJBBArB"
   },
   "outputs": [],
   "source": [
    "c[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8U6YUtU2BArE"
   },
   "source": [
    "Now that we have a `SkyCoord` object with the Gaia data, we can transform to other coordinate systems. For example, we can transform to the `Galactic` coordinate system (centered on the Sun but with the zero latitude approximately aligned with the Galactic plane) using the `.galactic` attribute (this works for any of the [built-in Astropy coordinate frames](http://docs.astropy.org/en/latest/coordinates/index.html#reference-api), e.g., `.fk5` should also work):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "31oHRJtqBArF"
   },
   "outputs": [],
   "source": [
    "c.galactic[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "K-yeGnD5BArI"
   },
   "source": [
    "The `Galactic` frame is still centered on the solar system barycenter, whereas we want to compute the positions and velocities of our sample of stars in a Galactocentric frame, centered on the center of the Milky Way. To do this transformation, Astropy provides the `Galactocentric` frame class, which allows us to use our own conventions for, e.g., the distance from the sun to the Galactic center (`galcen_distance`) or the height of the Sun over the Galactic midplane (`z_sun`). Let's look at the default values for the solar position and velocity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5Cmn3hp3BArI"
   },
   "outputs": [],
   "source": [
    "coord.Galactocentric()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GQYgwepOBArM"
   },
   "source": [
    "We'll instead use a distance of 8.1 kpc — more consistent with the [recent results from the GRAVITY collaboration](https://arxiv.org/abs/1807.09409) — and a solar height of 0 pc. We'll use the default solar velocity (see output above). We can transform our data to this frame using the `transform_to()` method by specifying the `Galactocentric` frame with our adopted values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zv9EtabKBArN"
   },
   "outputs": [],
   "source": [
    "galcen = c.transform_to(coord.Galactocentric(z_sun=0*u.pc, \n",
    "                                             galcen_distance=8.1*u.kpc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vH3XYYcZBArP"
   },
   "source": [
    "The `galcen` object now contains the data for our sample, but in the Galactocentric frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uOxijJC7BArQ"
   },
   "outputs": [],
   "source": [
    "galcen[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "It3RwYE1BArW"
   },
   "source": [
    "We can access the positions of the stars using the `.x`, `.y`, and `.z` attributes, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1S0jYb7ABArX"
   },
   "outputs": [],
   "source": [
    "plt.hist(galcen.z.value, bins=np.linspace(-110, 110, 32))\n",
    "plt.xlabel('$z$ [{0:latex_inline}]'.format(galcen.z.unit));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-y5-dreWBArc"
   },
   "source": [
    "Similarly, for the velocity components, we can use `.v_x`, `.v_y`, and `.v_z`. For example, to create a classic \"UV\" plane velocity plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "txcRdwL2BArd"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n",
    "\n",
    "ax.plot(galcen.v_x.value, galcen.v_y.value,\n",
    "        marker='.', linestyle='none', alpha=0.5)\n",
    "\n",
    "ax.set_xlim(-125, 125)\n",
    "ax.set_ylim(200-125, 200+125)\n",
    "\n",
    "ax.set_xlabel('$v_x$ [{0:latex_inline}]'.format(u.km/u.s))\n",
    "ax.set_ylabel('$v_y$ [{0:latex_inline}]'.format(u.km/u.s))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CFIdQARSBArh"
   },
   "source": [
    "Along with astrometric and radial velocity data, `Gaia` also provides photometric data for three photometric bandpasses: the broad-band `G`, the blue `BP`, and the red `RP` magnitudes. Let's make a Gaia color-magnitude diagram using the $G_{\\rm BP}-G_{\\rm RP}$ color and the absolute $G$-band magnitude $M_G$. We'll compute the absolute magnitude using the distances we computed earlier — Astropy `Distance` objects have a convenient `.distmod` attribute that provides the distance modulus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7SOj_ps_BArj"
   },
   "outputs": [],
   "source": [
    "M_G = gaia_data['phot_g_mean_mag'] - dist.distmod\n",
    "BP_RP = gaia_data['phot_bp_mean_mag'] - gaia_data['phot_rp_mean_mag']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "06YlJjNXBAro"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n",
    "\n",
    "ax.plot(BP_RP, M_G, \n",
    "        marker='.', linestyle='none', alpha=0.3)\n",
    "\n",
    "ax.set_xlim(0, 3)\n",
    "ax.set_ylim(11, 1)\n",
    "\n",
    "ax.set_xlabel('$G_{BP}-G_{RP}$')\n",
    "ax.set_ylabel('$M_G$')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iSoc8WLuBArt"
   },
   "source": [
    "In the above, there is a wide range of main sequence star masses which have a range of lifetimes. The most massive stars were likely born in the thin disk and their orbits therefore likely have smaller vertical amplitudes than the typical old main sequence star. To compare, we'll create two sub-selections of the Gaia CMD to select massive and low-mass main sequence stars from the CMD for comparison. You may see two ``RuntimeWarning``(s) from running the next cell — these are expected and it's safe to ignore them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6b4u3ROdBAru"
   },
   "outputs": [],
   "source": [
    "np.seterr(invalid=\"ignore\")\n",
    "hi_mass_mask = ((BP_RP > 0.5*u.mag) & (BP_RP < 0.7*u.mag) & \n",
    "                (M_G > 2*u.mag) & (M_G < 3.75*u.mag) & \n",
    "                (np.abs(galcen.v_y - 220*u.km/u.s) < 50*u.km/u.s))\n",
    "\n",
    "lo_mass_mask = ((BP_RP > 2*u.mag) & (BP_RP < 2.4*u.mag) & \n",
    "                (M_G > 8.2*u.mag) & (M_G < 9.7*u.mag) &\n",
    "                (np.abs(galcen.v_y - 220*u.km/u.s) < 50*u.km/u.s))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QPdWEdS2P0xT"
   },
   "source": [
    "Let's also define default colors to use when visualizing the high- and low-mass stars:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xNa7eJsXP3PY"
   },
   "outputs": [],
   "source": [
    "hi_mass_color = 'tab:red'\n",
    "lo_mass_color = 'tab:purple'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PtvmRVJsBArx"
   },
   "source": [
    "Let's now visualize these two CMD selections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4GkSKHNmBArx"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n",
    "\n",
    "ax.plot(BP_RP, M_G, \n",
    "        marker='.', linestyle='none', alpha=0.1)\n",
    "\n",
    "for mask, color in zip([lo_mass_mask, hi_mass_mask],\n",
    "                       [lo_mass_color, hi_mass_color]):\n",
    "    ax.plot(BP_RP[mask], M_G[mask], \n",
    "            marker='.', linestyle='none', \n",
    "            alpha=0.5, color=color)\n",
    "\n",
    "ax.set_xlim(0, 3)\n",
    "ax.set_ylim(11, 1)\n",
    "\n",
    "ax.set_xlabel('$G_{BP}-G_{RP}$')\n",
    "ax.set_ylabel('$M_G$')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sIt0gkXiBAr2"
   },
   "source": [
    "Thus far, we've used the color-magnitude diagram (using parallaxes and photometry from Gaia to compute absolute magnitudes) to select samples of high- and low-mass stars based on their colors.\n",
    "\n",
    "In what follows, we'll compute Galactic orbits for stars in the high- and low-mass star selections above and compare."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uuHWJMojBAr3"
   },
   "source": [
    "## Using `gala` to numerically integrate Galactic stellar orbits\n",
    "\n",
    "`gala` is an Astropy affiliated package for Galactic dynamics. `gala` provides functionality for representing analytic mass models that are commonly used in Galactic dynamics contexts for numerically integrating stellar orbits. For examples, see Chapter 3 of Binney and Tremaine (2008). The gravitational potential models are defined by specifying parameters like mass, scale radii, or shape parameters and can be combined. Once defined, they can be used in combination with numerical integrators provided in `gala` to compute orbits. `gala` comes with a pre-defined, multi-component, but [simple model for the Milky Way](http://gala.adrian.pw/en/latest/potential/define-milky-way-model.html) that can be used for orbit integrations. Let's create an instance of the `MilkyWayPotential` model and integrate orbits for the high- and low-mass main sequence stars selected above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VvhsMNVrBAr4"
   },
   "outputs": [],
   "source": [
    "milky_way = gp.MilkyWayPotential()\n",
    "milky_way"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aNiIVb5qBAr8"
   },
   "source": [
    "This model has mass components for the Galactic disk, bulge, nucleus, and halo, and the parameters were defined by fitting measurements of the Milky Way enclosed mass at various radii. See [this document](http://gala.adrian.pw/en/latest/potential/define-milky-way-model.html) for more details. The parameters of the `MilkyWayPotential` can be changed by passing in a dictionary of parameter values to argument names set by the component names. For example, to change the disk mass to make it slightly more massive (the choice `8e10` is arbitrary!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5c2y3aPQBAr9"
   },
   "outputs": [],
   "source": [
    "different_disk_potential = gp.MilkyWayPotential(disk=dict(m=8e10*u.Msun))\n",
    "different_disk_potential"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "05Mjz7wFBAr_"
   },
   "source": [
    "To integrate orbits, we have to combine the mass model with a reference frame into a `Hamiltonian` object. If no reference frame is passed in, it's assumed that we are in a static inertial frame moving with the center of the mass model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ey_b5jgHBAsA"
   },
   "outputs": [],
   "source": [
    "H = gp.Hamiltonian(milky_way)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6_PN-xWcBAsD"
   },
   "source": [
    "Now that we have the mass model, we can integrate orbits. Let's now define initial conditions for subsets of the high- and low-mass star selections we did above. Initial conditions in `gala` are specified by creating `PhaseSpacePosition` objects. We can create these objects directly from a `Galactocentric` object, like we have defined above from transforming the Gaia data — we first have to extract the data with a Cartesian representation. We can do this by calling `galcen.cartesian`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TUw2WHDCBAsF"
   },
   "outputs": [],
   "source": [
    "w0_hi = gd.PhaseSpacePosition(galcen[hi_mass_mask].cartesian)\n",
    "w0_lo = gd.PhaseSpacePosition(galcen[lo_mass_mask].cartesian)\n",
    "w0_hi.shape, w0_lo.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "DAY6mDHEBAsI"
   },
   "source": [
    "From the above, we can see that we have 185 high-mass star and 577 low-mass stars in our selections. To integrate orbits, we call the `.integrate_orbit()` method on the Hamiltonian object we defined above, and pass in initial conditions. We also have to specify the timestep for integration, and how long we want to integrate for. We can do this by either specifying the amount of time to integrate for, or by specifying the number of timesteps. Let's specify a timestep of 1 Myr and a time of 500 Myr (approximately two revolutions around the Galaxy for a Sun-like orbit):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hauyeC8aBAsK"
   },
   "outputs": [],
   "source": [
    "orbits_hi = H.integrate_orbit(w0_hi, dt=1*u.Myr, \n",
    "                              t1=0*u.Myr, t2=500*u.Myr)\n",
    "\n",
    "orbits_lo = H.integrate_orbit(w0_lo, dt=1*u.Myr, \n",
    "                              t1=0*u.Myr, t2=500*u.Myr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HsCFYzdGBAsL"
   },
   "source": [
    "By default this uses a [Leapfrog](https://en.wikipedia.org/wiki/Leapfrog_integration) numerical integration scheme, but the integrator can be customized — see the `gala` [examples](http://gala.adrian.pw/en/latest/examples/integrate-potential-example.html) for more details.\n",
    "\n",
    "With the orbit objects in hand, we can continue our comparison of the orbits of high-mass and low-mass main sequence stars in the solar neighborhood. Let's start by plotting a few orbits. The `.plot()` convenience function provides a quick way to visualize orbits in three Cartesian projections. For example, let's plot the first orbit in each subsample on the same figure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7vF1SKOZBAsM"
   },
   "outputs": [],
   "source": [
    "fig = orbits_hi[:, 0].plot(color=hi_mass_color)\n",
    "_ = orbits_lo[:, 0].plot(axes=fig.axes, color=lo_mass_color)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cTpara5NN482"
   },
   "source": [
    "Note in the above figure that the orbits are almost constrained to the x-y plane: the excursions are much larger in the x and y directions as compared to the z direction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3O8_RipDBAsS"
   },
   "source": [
    "The default plots show all Cartesian projections. This can be customized to, for example, only show specified components (including velocity components):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XjCNGo-9BAsT"
   },
   "outputs": [],
   "source": [
    "fig = orbits_hi[:, 0].plot(['x', 'v_x'], \n",
    "                           auto_aspect=False, \n",
    "                           color=hi_mass_color)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HkKgLLB6BAsW"
   },
   "source": [
    "The representation can also be changed, for example, to a cylindrical representation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zBp9vI1EBAsW"
   },
   "outputs": [],
   "source": [
    "fig = orbits_hi[:, 0].cylindrical.plot(['rho', 'z'], \n",
    "                                       color=hi_mass_color,\n",
    "                                       label='high mass')\n",
    "_ = orbits_lo[:, 0].cylindrical.plot(['rho', 'z'], color=lo_mass_color,\n",
    "                                     axes=fig.axes,\n",
    "                                     label='low mass')\n",
    "\n",
    "fig.axes[0].legend(loc='upper left')\n",
    "fig.axes[0].set_ylim(-0.3, 0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "StVu7BGYSBcn"
   },
   "source": [
    "Already in the above plot we can see that the high-mass star has an orbit with smaller eccentricity (smaller radial variations) and smaller vertical oscillations as compared to the low-mass star. Below, we'll quantify this and look at the vertical excursions of all of the high- and low-mass stars, respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WVzt3tfDBAsY"
   },
   "source": [
    "Let's now compare the vertical amplitudes of the orbits in each of our sub-selections! We can compute the (approximate) maximum vertical height of each orbit using the convenience method `.zmax()` (you can see a list of all convenience methods on the `Orbit` object [in the Gala documentation here](http://gala.adrian.pw/en/latest/api/gala.dynamics.Orbit.html#gala.dynamics.Orbit)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rXrbqbwwBAsb"
   },
   "outputs": [],
   "source": [
    "zmax_hi = orbits_hi.zmax(approximate=True)\n",
    "zmax_lo = orbits_lo.zmax(approximate=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Z-f_FsitBAsm"
   },
   "source": [
    "Let's make histograms of the maximum $z$ heights for these two samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dw4EDSwVBAsn"
   },
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 2, 50)\n",
    "\n",
    "plt.hist(zmax_hi.value, bins=bins, \n",
    "         alpha=0.4, density=True, label='high-mass', \n",
    "         color=hi_mass_color)\n",
    "plt.hist(zmax_lo.value, bins=bins, \n",
    "         alpha=0.4, density=True, label='low-mass',\n",
    "         color=lo_mass_color);\n",
    "\n",
    "plt.legend(loc='best', fontsize=14)\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel(r\"$z_{\\rm max}$\" + \" [{0:latex}]\".format(zmax_hi.unit))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-_CyXGAxBAsr"
   },
   "source": [
    "The distribution of $z$-heights for the low-mass (i.e. typically older) stars is more extended, as we predicted!\n",
    "\n",
    "In this tutorial, we've used `astroquery` to query the Gaia science archive to retrieve kinematic and photometric data for a small sample of stars with well-measured parallaxes from Gaia DR2. We used the colors and absolute magnitudes of these stars to select subsamples of high- and low-mass stars, which, on average, will provide us with subsamples of stars that are younger and older, respectively. We then constructed a model for the gravitational field of the Milky Way and numerically integrated the orbits of all stars in each of the two subsamples. Finally, we used the orbits to compute the maximum height that each star reaches above the Galactic midplane and showed that the younger (higher-mass) stars tend to have smaller excursions from the Galactic plane, consistent with the idea that stars are either born in a \"thinner\" disk and dynamically \"heated,\" or that older stars formed with a larger vertical scale-height."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "52uk3tvIKXdn"
   },
   "source": [
    "## Exercises\n",
    "\n",
    "1. Some of the low-mass star orbits have large vertical excursions from the Galactic disk (up to and above 1.5 kpc) and could therefore be stellar halo stars rather than part of the Galactic disk. Use the zmax values to select a few of these stars and plot their full orbits. Do these stars look like they are part of the disk? Why / why not?\n",
    "2. [Orbit](http://gala.adrian.pw/en/latest/dynamics/orbits-in-detail.html) objects also provide methods for computing apocenter and pericenter distances and eccentricities. Which types of stars (high-mass or low-mass) tend to have high eccentricity orbits within the Galaxy? Similar to the plot above, make a plot showing the two distributions of eccentricity values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "i_D4XzdeVEnE"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "gaia-galactic-orbits.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
